#ifndef BLOCK_EPILOGUE_SINGLE_OP_FP_H
#define BLOCK_EPILOGUE_SINGLE_OP_FP_H

#include "./block_epilogue_single_op.h"
#include "../utils/constant_values.h"
#include "../utils/sync_utils.h"

template<typename InputDataType_, typename OutputDataType_, typename FusionOp_>
class BlockEpilogueSingleOp<InputDataType_, OutputDataType_, FusionOp_,
    AscendC::Std::enable_if_t<
        AscendC::Std::is_same_v<InputDataType_, float> && AscendC::Std::is_same_v<OutputDataType_, bfloat16_t>
    >
> {
public:
    using InputDataType = InputDataType_;
    using OutputDataType = OutputDataType_;
    using FusionOp = FusionOp_;

    using FusionArguments = typename FusionOp::Arguments;
    using FusionParams = typename FusionOp::Params;

    struct Arguments {
        int32_t dataCount{0};
        FusionArguments fusionOp1Args{};
        Arguments() = default;
    };

    struct Params {
        int32_t dataCount{0};
        FusionParams fusionOp1Params{};
        Params() = default;
    };

private:
    AscendC::TBuf<> paraBuf_;
    AscendC::TBuf<> inputBuf_;
    AscendC::TBuf<> resBuf_;
    AscendC::LocalTensor<InputDataType> paraLocal_;
    AscendC::LocalTensor<InputDataType> inputLocal_;
    AscendC::LocalTensor<OutputDataType> resLocal_;
    FusionOp op;

    struct VecSyncs {
        SEvent<AscendC::HardEvent::MTE2_V> mte2ToV;
        SEvent<AscendC::HardEvent::V_MTE3> vToMTE3;
        SEvent<AscendC::HardEvent::MTE3_MTE2> mte3ToMTE2;
    };

public:
    __host_aicore__ static Params InitParams(Arguments const& args) {
        return {
            args.dataCount,
            FusionOp::InitParams(args.fusionArgs)
        };
    }

    __host_aicore__ static size_t GetWorkSpaceSize(Arguments const& args) {
        return 0;
    }

    __aicore__ inline void Init(Params const& params) {
        AscendC::TPipe* pipe = GetTPipePtr();
        pipe->Reset();
        pipe->InitBuffer(paraBuf_, UB_SIZE / 3);
        pipe->InitBuffer(inputBuf_, UB_SIZE / 3);
        pipe->InitBuffer(resBuf_, UB_SIZE / 3);

        paraLocal_  = paraBuf_.template Get<InputDataType>();
        inputLocal_ = inputBuf_.template Get<InputDataType>();
        resLocal_   = resBuf_.template Get<OutputDataType>();
    }

    __aicore__ inline void Run(
        const AscendC::GlobalTensor<OutputDataType>& outputGm,
        const AscendC::GlobalTensor<InputDataType>& inputGm,
        const AscendC::GlobalTensor<InputDataType>& paraGm,
        Params const& params
    ) {
        int32_t vecIdx = AscendC::GetSubBlockIdx();
        constexpr int32_t chunkSize = UB_SIZE / 3/ sizeof(float);
        constexpr int32_t policyTwoChunkCheck = chunkSize * 2;
        constexpr int32_t policyOneChunkCheck = chunkSize;

        VecSyncs vss;
        const int32_t dataCount = params.dataCount;
        for (int32_t offset = 0; offset < dataCount;) {
            int32_t vecOffset{0};
            int32_t dataCountForProcess{0};

            if (dataCount >= policyTwoChunkCheck) {
                vecOffset = offset + vecIdx * chunkSize;
                dataCountForProcess = chunkSize;
                offset += 2 * chunkSize;
            } else if (dataCount > policyOneChunkCheck) {
                vecOffset = offset + vecIdx * chunkSize;
                dataCountForProcess = (vecIdx == 0 ? chunkSize : dataCount - offset - chunkSize);
                offset = dataCount;
            } else {
                if (vecIdx == 1) {
                    break;
                }
                vecOffset = offset;
                dataCountForProcess = dataCount - offset;
                offset = dataCount;
            }

            ProcessChunk(
                outputGm[vecOffset],
                inputGm[vecOffset],
                paraGm[vecOffset],
                dataCountForProcess,
                vss
            );
        }
    }
private:
    __aicore__ inline void ProcessChunk(
        const AscendC::GlobalTensor<OutputDataType>& outputGm,
        const AscendC::GlobalTensor<InputDataType>& inputGm,
        const AscendC::GlobalTensor<InputDataType>& paraGm,
        int32_t chunkSize,
        VecSyncs& vecSyncs
    ) {
        AscendC::DataCopy(inputLocal_, inputGm, chunkSize);
        AscendC::DataCopy(paraLocal_, paraGm, chunkSize);
        vecSyncs.mte2ToV.SetFlag();
        vecSyncs.mte2ToV.WaitFlag();

        op.Run(resLocal_, paraLocal_, inputLocal_, chunkSize);

        vecSyncs.vToMTE3.SetFlag();
        vecSyncs.vToMTE3.WaitFlag();
        AscendC::DataCopy(outputGm, resLocal_, chunkSize);
        vecSyncs.mte3ToMTE2.SetFlag();
        vecSyncs.mte3ToMTE2.WaitFlag();
    }
};

#endif // !BLOCK_EPILOGUE_SINGLE_OP_BF_H